// if compiling for Pentium 4 or later uncomment line below
// #define SSE2 
#define CHECK_ALL
#include "fermiqcd.h"
#include "fermiqcd_fermi_smearing.h"
#include "fermiqcd_JIM_IO.h"

void apply_gamma_dot_D(fermi_field& psi, gauge_field& U) {
  fermi_field psi_rot(psi.lattice(),psi.nspin,psi.nc);
  psi_rot=0;
  mdp_site x(psi.lattice());
  forallsites(x)
    for(int mu=1;mu<U.ndim;mu++)
      for(int a=0;a<psi.nspin;a++)
        for(int b=0;b<psi.nspin;b++) {
          psi(x,a)+=0.5*Gamma[mu](a,b)*(U(x,mu)*psi_rot(x+mu,b)-(U(x,-1,mu))*psi_rot(x-mu,b));
        }
  psi.update(); 
}

// For staggered->naive propagator conversion
// Here only for a check
mdp_matrix OmegaMatrix(site x) { 
  mdp_matrix M(4,4); 
  M=1; 
  if(x(3)%2==1) M=Gamma[3]*M; 
  if(x(2)%2==1) M=Gamma[2]*M; 
  if(x(1)%2==1) M=Gamma[1]*M; 
  if(x(0)%2==1) M=Gamma[0]*M;  
  return M;
} 

// this is actually used and compared with function above but faster and 
// saves memory
mdp_complex Omega(site x, int a, int b) {
  static bool done=false;
  static mdp_matrix m[16];
  static mdp_matrix tmp;
  int i;
  if(!done) {
    for(i=0; i<16; i++) {
      tmp=mdp_identity(4);
      if((i>>3) & 0x1) tmp=Gamma[3]*tmp; 
      if((i>>2) & 0x1) tmp=Gamma[2]*tmp; 
      if((i>>1) & 0x1) tmp=Gamma[1]*tmp; 
      if((i>>0) & 0x1) tmp=Gamma[0]*tmp;        
      m[i]=tmp;
    }
    done=true;
  }  
  i=x(0)%2 + 2*(x(1)%2 + 2*(x(2)%2 + 2*(x(3)%2)));
#ifdef CHECK_ALL
  for(int c=0; c<4; c++)
    for(int d=0; d<4; d++)
      if(m[i](a,b)!=OmegaMatrix(x)(a,b)) cout << "ERRORRRRRRRRRRRRRRRRRR!";
#endif
  return m[i](a,b);
}

// this maps a staggered prop chi (as created by Jim) 
// into a naive prop using the Omega function above
inline mdp_complex S_light(mdp_matrix_field &chi, mdp_site &x, 
			   int b, int j, 
			   int a, int i) {
  // This is in accord to Todd's stag2naive.cc code
  return Omega(x,b,a)*chi(x,i,j); 
}

class open_indices {
public:
  mdp_complex c[4][4][3][3];
};

void empty(mdp_field<open_indices> &open_prop) {
  mdp_site x1d(open_prop.lattice());
  int a,b,i,j;
  forallsites(x1d) 
    for(a=0; a<4; a++)
      for(b=0; b<4; b++)
	for(i=0; i<3; i++)
	  for(j=0; j<3; j++)
	    open_prop(x1d).c[a][b][i][j]=0;
}

int main(int argc, char** argv) {
  mdp.open_wormholes(argc, argv);
  define_base_matrices("FERMILAB");

  if(argc<6) {
    mdp << "Usage:\n";
    mdp << argv[0] 
	<< " 96x40x40x40 gaugefile staggered_prop_filename kappa_heavy csw_heavy [ap] [rp] [d1]\n"
	<< "[ap]=20 for 10^-20 absolute precision\n"
	<< "[rp]=12 for 12^-12 relative precision\n";
    exit(1);
  }

  if(!is_file(argv[1])) {
    mdp << "Unbale to open gauge configuration\n";
    exit(1);
  }
  // read the lattice size from the gauge configuration itself
  int L[4];
  scanf(argv[1],"%ix%ix%ix%i", L,L+1,L+2,L+3);
  int nc=3; // number of colors;
  mdp_lattice lattice(4,L,default_partitioning<0>,torus_topology,0,1,false);
  mdp_lattice lattice1d(1,L,default_partitioning<0>,torus_topology,0,1,false);
  gauge_field U(lattice,nc);  
  fermi_field source(lattice,nc);
  fermi_field sink_heavy(lattice,nc); // light quark propagator
  mdp_matrix_field chi(lattice,nc,nc);
  mdp_site x(lattice); // site on the 4d lattice
  mdp_site x1d(lattice1d); // site on the projected lattice
  int a,b,c,d,i,j,k; // temp variables
  mdp_complex coeff;

  mdp_field<open_indices> open_prop(lattice1d); // field for output open prop

  char filename[1024]; // UNSAFE but OK
  char output_filename[1024]; // UNSAFE but OK
  string gauge_filename;
  string staggered_prop_filename;
  coefficients heavy_quark;
  coefficients smearing_coefficients;

  // making these numbers up
  gauge_filename=argv[2];
  staggered_prop_filename=argv[3];  
  heavy_quark["kappa"]=atof(argv[4]); 
  heavy_quark["c_{SW}"]=atof(argv[5]); 
  float ap=1e-20; // inversion absolute precision
  if(argc>6) ap=pow(10.0,-atoi(argv[6]));
  float rp=1e-12; // inversion relative precision
  if(argc>7) rp=pow(10.0,-atoi(argv[7]));

  smearing_coefficients["factor"]=0.5; // FIX THIS
  smearing_coefficients["steps"]=10; // FIX THIS

  mdp << "Target absolute precision = " << ap << endl;
  mdp << "Target relative precision = " << rp << endl;

  // load gauge configuration
  //Use MILC config in JIM_IO format
  //read in gauge config, associate gauge config with U fields
  MetaHash UserMeta;
  JIM_load<float> ( U, gauge_filename.c_str(), UserMeta );
  compute_em_field(U);

#ifdef SSE2 
  default_fermi_action=FermiCloverActionSSE2::mul_Q;
#endif

  //Load Staggered Prop
  MetaHash StagMeta;
  JIM_load<float> (chi,staggered_prop_filename.c_str(),StagMeta);

  for(int smearing=0; smearing<2; smearing++) {
    for(int rotation=0; rotation<2; rotation++) {

      empty(open_prop);
    
      for(b=0; b<4; b++)
	for(j=0; j<3; j++) {
	  if(!rotation) {	  
	    sprintf(filename,"%s_heavy_%s_%s_%i_%i", argv[2],argv[4],argv[5],b,j);
	  } else {
	    sprintf(filename,"%s_heavy_%s_%s_%i_%i_rot",argv[2],argv[4],argv[5],b,j);
	  }	
	  if(smearing==0) {
	    // make heavy prop component
	    source=0;
	    x.set(0);
	    source(x,b,j)=1;
	    if(rotation) apply_gamma_dot_D(source,U);
	    mul_invQ(sink_heavy,source,U,heavy_quark,ap,rp,5000);      
	    sink_heavy.save(filename);
	  } else {
	    sink_heavy.load(filename);
	    WupperthalSmearing::smear(sink_heavy,U,smearing_coefficients);
	  }
	  
	  d=Gamma5_idx[b];
	  coeff=Gamma5_val[b];
	  forallsites(x) {
	    x1d.set(x(0)); // project a 4d site x into a 1d site x1d
	    for(a=0; a<4; a++)
	      for(c=0; c<4; c++)
		for(i=0; i<3; i++)
		  for(k=0; k<3; k++) {
		    /*
		      Accordint to Todd's Email
		      O_{x,d,a,j,i}=Gamma5_{db}*conj(S_h{x,c,b,k,j))*S_l(x,c,a,k,i)
		      And I agree
		      
		      Jim suggests we also open spin indices at sink
		  */
		  open_prop(x1d).c[a][d][i][j]+=
		    coeff*conj(sink_heavy(x,c,k))*S_light(chi,x,c,a,k,i);
		}
	}
      }
      sprintf(output_filename,"%s_light_%s_%s_heavy_%s_%s_%i_%i_open",
	      argv[1], argv[2], argv[3], argv[4], argv[5],rotation,smearing);
      open_prop.save(output_filename);
    }
  }
  mdp.close_wormholes();
  return 0;
}

